from .image_convert_base import ConvertBase
import cv2 as cv
import numpy as np



class ScaleConvert(ConvertBase):
    
    def __init__(self, use_rate : float = 0.5, min_scale : float = 0.4, max_scale : float = 1, paddle_color : list = [0, 0, 0]):
        super().__init__(use_rate=use_rate)
        
        self.min_scale = min_scale
        self.max_scale = max_scale    
        self.paddle_color = paddle_color
        
    
    def __get_scale(self):
        return np.random.randint(self.min_scale * 10, self.max_scale * 10) / 10.0
    
    
    def convert(self, img, boxes, points):
        h, w, _ = img.shape
        target_img = np.zeros(img.shape, np.uint8)
        target_img[:, :] = self.paddle_color
        
        scale = self.__get_scale()
        target_size = (int(w * scale), int(h * scale))
        img = cv.resize(img, target_size)
        
        x_offset = int((w -  target_size[0]) / 2)
        y_offset = int((h - target_size[1]) / 2)
        # print(target_size, w, h, x_offset, y_offset, img.shape, scale)
        target_img[y_offset:y_offset + target_size[1], x_offset:x_offset + target_size[0]] = img
        
        if boxes is not None:
            boxes *= scale
            boxes[:, [0, 2]] += x_offset
            boxes[:, [1, 3]] += y_offset
        
        if points is not None:
            points *= scale
            points[:, 0] += x_offset
            points[:, 1] += y_offset
            
        
        return target_img, boxes, points
            
            
    